#include <ros/ros.h>
#include <geometry_msgs/PoseStamped.h>
#include <geometry_msgs/TransformStamped.h>
#include <sensor_msgs/Imu.h>
#include <Eigen/Eigen>
#include <Eigen/Core>
#include <Eigen/Geometry>
#include <fstream>
#include "loc_correction/filter_utils.h"
#include "loc_correction/real_time_outlier_removal.h"
#include "loc_correction/filter2.h"
#include "loc_correction/optimizecampos.h"
// #include "loc_correction/math_utils.h"
// #include "loc_correction/kalman_filter.h"
// pos_uav_init: 
//     [ 0.1502688229084015,
//       0.222562775015831,
//       0.1135887014865875]

// pos_uav2uwb01_init: 
//     [ 0.9897,
//       0.6753,
//       -0.01955]

// pos_uwb01_init: 
//     [ −0.839431177,
//       −0.452737225,
//       0.133138701]




class cam_pos_counter{
public:

    cam_pos_counter() : 
    cam_pos_filter(20), cam_orientation_filter(20),
    outlier_removal_x(), outlier_removal_y(), outlier_removal_z(),
    orientation_outlier_removal_w(), orientation_outlier_removal_x(), orientation_outlier_removal_y(), orientation_outlier_removal_z()    
    {                                                  
        sub_uav_ = nh_.subscribe("/uav201/mocap/pos", 1, &cam_pos_counter::uav_cb, this);//around 90hz
        sub_imu_ = nh_.subscribe("/IMU_data", 1, &cam_pos_counter::imu_cb, this);//around 100hz
        sub_pnp_ = nh_.subscribe("/camA/single_cam_process_ros/ir_mono/T_base_to_estimation", 1, &cam_pos_counter::pnp_cb, this);//around 30hz
        
        // imu orientation
        pub_cam_ = nh_.advertise<geometry_msgs::PoseStamped>("/cam_pos", 1);
        filename = "/home/jigin/multi_loop/src/loc_correction/data/cam_pos.txt";
        uav_mat = Eigen::Matrix4d::Identity();
        pnp_mat = Eigen::Matrix4d::Identity();
        cam_mat = Eigen::Matrix4d::Identity();
        ref_uav_pos = Eigen::Matrix4d::Identity();

        uav_flag = false;
        pnp_flag = false;
        imu_flag = false;

        // 清空当前.txt文件
        std::ofstream file;
        file.open(filename, std::ios::out | std::ios::trunc);
        if (file.is_open()) {
            file.close();
        } else {
            ROS_ERROR("Unable to open file: %s", filename.c_str());
        }
    }

    void uav_cb(const geometry_msgs::PoseStamped& msg) {
        // ROS_INFO("Received uav pos : [%.2f , %.2f , %.2f]", msg.pose.position.x, msg.pose.position.y, msg.pose.position.z);
        // Eigen::Vector3d uav_raw_position(msg.pose.position.x, msg.pose.position.y, msg.pose.position.z);
        // Eigen::Vector3d uav_smooth_position = uav_filter.apply(uav_raw_position);
        // uav_mat.block<3, 1>(0, 3) << uav_smooth_position.x(), uav_smooth_position.y(), uav_smooth_position.z();
        uav_mat.block<3, 1>(0, 3) << msg.pose.position.x, msg.pose.position.y, msg.pose.position.z;

        // Eigen::Quaterniond q(msg.pose.orientation.w, msg.pose.orientation.x, msg.pose.orientation.y, msg.pose.orientation.z);
        // uav_mat.block<3, 3>(0, 0) = q.toRotationMatrix();
        uav_flag = true;
    }

    void imu_cb(const sensor_msgs::Imu& msg) {
        // ROS_INFO("Received imu data : [%.2f , %.2f , %.2f , %.2f]", msg.orientation.w, msg.orientation.x, msg.orientation.y, msg.orientation.z);
        Eigen::Quaterniond q(msg.orientation.w, msg.orientation.x, msg.orientation.y, msg.orientation.z);
        uav_mat.block<3, 3>(0, 0) = q.toRotationMatrix();
        imu_flag = true;
    }


    void pnp_cb(const geometry_msgs::TransformStamped& msg) {
        // ROS_INFO("Received  : [%.2f , %.2f , %.2f ]", msg.transform.translation.x, msg.transform.translation.y, msg.transform.translation.z);

        // Eigen::Vector3d raw_position(msg.transform.translation.x, msg.transform.translation.y, msg.transform.translation.z);
        // Eigen::Vector3d smooth_position = pnp_filter.apply(raw_position);
        // pnp_mat.block<3, 1>(0, 3) << smooth_position.x(), smooth_position.y(), smooth_position.z();

        // double smooth_z = kf_x.update(msg.transform.translation.z);

        pnp_mat.block<3, 1>(0, 3) << msg.transform.translation.x, msg.transform.translation.y, msg.transform.translation.z;

        Eigen::Quaterniond q(msg.transform.rotation.w, msg.transform.rotation.x, msg.transform.rotation.y, msg.transform.rotation.z);
        pnp_mat.block<3, 3>(0, 0) = q.toRotationMatrix();
        pnp_flag = true;
        if(uav_flag && pnp_flag && imu_flag) {
            // std::cout<<pnp_mat<<std::endl;
            cam_mat = pnp_mat.inverse() * uav_mat;
            // std::cout<<cam_mat<<std::endl;

            //remove outlier==============================
            double filtered_pos_x = outlier_removal_x.process(cam_mat(0,3));
            double filtered_pos_y = outlier_removal_y.process(cam_mat(1,3));
            double filtered_pos_z = outlier_removal_z.process(cam_mat(2,3));
            //============================================

            //using kalman filter=========================
            // double smooth_x = kf_x.update(cam_mat(0, 3));
            // double smooth_y = kf_y.update(cam_mat(1, 3));
            // double smooth_z = kf_z.update(cam_mat(2, 3));
            // cam_msg.pose.position.x = smooth_x;
            // cam_msg.pose.position.y = smooth_y;
            // cam_msg.pose.position.z = smooth_z;
            //============================================

            //using filter_utils==========================
            Eigen::Vector3d raw_position(filtered_pos_x, filtered_pos_y, filtered_pos_z);
            Eigen::Vector3d smooth_position = cam_pos_filter.apply(raw_position);
            cam_msg.pose.position.x = smooth_position.x();
            cam_msg.pose.position.y = smooth_position.y();
            cam_msg.pose.position.z = smooth_position.z();
            //============================================

            //origin cam_pos==============================
            // cam_msg.pose.position.x = filtered_pos_x;
            // cam_msg.pose.position.y = filtered_pos_y;
            // cam_msg.pose.position.z = filtered_pos_z;
            //============================================

            cam_msg.header.stamp = ros::Time::now();
            cam_msg.header.frame_id = "world";

            Eigen::Quaterniond quaternion(cam_mat.block<3,3>(0,0));
            //remove outlier==============================
            double filtered_orientation_w = orientation_outlier_removal_w.process(quaternion.w());
            double filtered_orientation_x = orientation_outlier_removal_x.process(quaternion.x());
            double filtered_orientation_y = orientation_outlier_removal_y.process(quaternion.y());
            double filtered_orientation_z = orientation_outlier_removal_z.process(quaternion.z());
            //============================================
            //using filter_utils==========================
            Eigen::Quaterniond raw_orientation(filtered_orientation_w, filtered_orientation_x, filtered_orientation_y, filtered_orientation_z);
            raw_orientation.normalize();
            Eigen::Quaterniond smooth_position_vec4 = cam_orientation_filter.apply(raw_orientation);
            cam_msg.pose.orientation.x = smooth_position_vec4.x();
            cam_msg.pose.orientation.y = smooth_position_vec4.y();
            cam_msg.pose.orientation.z = smooth_position_vec4.z();
            cam_msg.pose.orientation.w = smooth_position_vec4.w();
            //===========================================
            //original orientation=======================
            // cam_msg.pose.orientation.w = filtered_orientation_w;
            // cam_msg.pose.orientation.x = filtered_orientation_x;
            // cam_msg.pose.orientation.y = filtered_orientation_y;
            // cam_msg.pose.orientation.z = filtered_orientation_z;
            //===========================================

            // pub_cam_.publish(cam_msg);//publish

            // counting variance===========================
            // pos_x.push_back(cam_msg.pose.position.x);
            // pos_y.push_back(cam_msg.pose.position.y);
            // pos_z.push_back(cam_msg.pose.position.z);
            // if (pos_z.size() >= 29){
            //     double variance_z = MathUtils::calculate_variance(pos_z);
            //     std::cout<<variance_z<<std::endl;
            //     pos_x.clear();
            //     pos_y.clear();
            //     pos_z.clear();
            // }
            //==============================================


            // // // 记录当前帧位置数据in txt
            // ROS_INFO("Cam_pos: (%.3f, %.3f, %.3f)", cam_mat(0, 3), cam_mat(1, 3), cam_mat(2, 3));
            // std::ofstream file;
            // file.open(filename, std::ios::app);
            // if(file.is_open()) {
            //     file << cam_mat(0, 3) << "    ";
            //     file << cam_mat(1, 3) << "    ";
            //     file << cam_mat(2, 3) << "    ";
            //     file << cam_msg.pose.orientation.w << "    ";
            //     file << cam_msg.pose.orientation.x << "    ";
            //     file << cam_msg.pose.orientation.y << "    ";
            //     file << cam_msg.pose.orientation.z << "    ";
            //     file << '\n';
            // }
            // else {
            //     ROS_ERROR("Unable to open file: %s", filename.c_str());
            // }

            // uav_flag = false;
            // pnp_flag = false;
            // imu_flag = false;
        }

        if (uav_flag && pnp_flag && imu_flag)
        {
            Eigen::Vector3d cam_pos = cam_mat.block<3, 1>(0, 3);
            Eigen::Quaterniond cam_orientation(cam_mat.block<3, 3>(0, 0));

            optimizeCamPose(cam_pos, cam_orientation, pnp_mat, ref_uav_pos, 50, 0.01);

            cam_msg.pose.position.x = cam_pos.x();
            cam_msg.pose.position.y = cam_pos.y();
            cam_msg.pose.position.z = cam_pos.z();
            cam_msg.pose.orientation.x = cam_orientation.x();
            cam_msg.pose.orientation.y = cam_orientation.y();
            cam_msg.pose.orientation.z = cam_orientation.z();
            cam_msg.pose.orientation.w = cam_orientation.w();

            pub_cam_.publish(cam_msg);

            // // 记录当前帧位置数据in txt
            ROS_INFO("Cam_pos: (%.3f, %.3f, %.3f)", cam_mat(0, 3), cam_mat(1, 3), cam_mat(2, 3));
            std::ofstream file;
            file.open(filename, std::ios::app);
            if(file.is_open()) {
                file << cam_pos.x() << "    ";
                file << cam_pos.y() << "    ";
                file << cam_pos.z() << "    ";
                file << cam_orientation.x() << "    ";
                file << cam_orientation.y() << "    ";
                file << cam_orientation.z() << "    ";
                file << cam_orientation.w() << "    ";
                file << '\n';
            }
            else {
                ROS_ERROR("Unable to open file: %s", filename.c_str());
            }

            uav_flag = false;
            pnp_flag = false;
            imu_flag = false;

            
        }
            
    }


private:
    ros::NodeHandle nh_;
    ros::Subscriber sub_uav_;
    ros::Subscriber sub_pnp_;
    ros::Subscriber sub_imu_;
    ros::Publisher pub_cam_;
    Eigen::Matrix4d uav_mat;
    Eigen::Matrix4d pnp_mat;
    Eigen::Matrix4d cam_mat;
    Eigen::Matrix4d ref_uav_pos;

    geometry_msgs::PoseStamped cam_msg;
    bool uav_flag;
    bool pnp_flag;
    bool imu_flag;
    std::string filename;


    //filter_utils========================
    // MovingAverageFilter pnp_filter;
    // MovingAverageFilter uav_filter;
    MovingAverageFilter cam_pos_filter;
    MovingAverageFilter_quat cam_orientation_filter;


    // std::vector<double> pos_x;
    // std::vector<double> pos_y;
    // std::vector<double> pos_z;

    //kalman filter====================
    // KalmanFilter kf_x;
    // KalmanFilter kf_y;
    // KalmanFilter kf_z;


    //remove outliers==================
    RealTimeOutlierRemoval outlier_removal_x;
    RealTimeOutlierRemoval outlier_removal_y;
    RealTimeOutlierRemoval outlier_removal_z;
    RealTimeOutlierRemoval orientation_outlier_removal_w;
    RealTimeOutlierRemoval orientation_outlier_removal_x;
    RealTimeOutlierRemoval orientation_outlier_removal_y;
    RealTimeOutlierRemoval orientation_outlier_removal_z;



};

int main(int argc, char* argv[]) {
    ros::init(argc, argv, "cam_pos_counter");
    std::cout<<"cam_pos_reader start..."<<std::endl;
    cam_pos_counter reader;
    ros::spin();
    return 0;
}